"""
Ollama
Ollama
"""

import requests
import json
import logging
from typing import Dict, Any, List, Optional, Generator
import time

class OllamaClient:
    """Ollama"""
    
    def __init__(self, base_url: str = "http://localhost:11434"):
        self.base_url = base_url.rstrip('/')
        self.logger = logging.getLogger(__name__)
        
    def is_available(self) -> bool:
        """Ollama"""
        try:
            response = requests.get(f"{self.base_url}/api/tags", timeout=5)
            return response.status_code == 200
        except Exception as e:
            self.logger.error(f"Ollama: {e}")
            return False
    
    def list_models(self) -> List[Dict[str, Any]]:
        """"""
        try:
            response = requests.get(f"{self.base_url}/api/tags")
            response.raise_for_status()
            data = response.json()
            return data.get('models', [])
        except Exception as e:
            self.logger.error(f": {e}")
            return []
    
    def model_exists(self, model_name: str) -> bool:
        """"""
        models = self.list_models()
        return any(model['name'] == model_name for model in models)
    
    def pull_model(self, model_name: str, progress_callback=None) -> bool:
        """"""
        try:
            self.logger.info(f": {model_name}")
            
            response = requests.post(
                f"{self.base_url}/api/pull",
                json={"name": model_name},
                stream=True
            )
            response.raise_for_status()
            
            for line in response.iter_lines():
                if line:
                    try:
                        data = json.loads(line)
                        if progress_callback and 'completed' in data and 'total' in data:
                            progress = int((data['completed'] / data['total']) * 100)
                            progress_callback(progress, f": {model_name}")
                        
                        if data.get('status') == 'success':
                            self.logger.info(f": {model_name}")
                            return True
                            
                    except json.JSONDecodeError:
                        continue
            
            return True
            
        except Exception as e:
            self.logger.error(f": {e}")
            return False
    
    def generate(self, model: str, prompt: str, system: str = None, 
                temperature: float = 0.7, max_tokens: int = 2000,
                stream: bool = False) -> Dict[str, Any]:
        """"""
        try:
            payload = {
                "model": model,
                "prompt": prompt,
                "stream": stream,
                "options": {
                    "temperature": temperature,
                    "num_predict": max_tokens
                }
            }
            
            if system:
                payload["system"] = system
            
            response = requests.post(
                f"{self.base_url}/api/generate",
                json=payload,
                timeout=300  # 5
            )
            response.raise_for_status()
            
            if stream:
                return self._handle_stream_response(response)
            else:
                return response.json()
                
        except Exception as e:
            self.logger.error(f": {e}")
            return {"error": str(e)}
    
    def generate_stream(self, model: str, prompt: str, system: str = None,
                       temperature: float = 0.7, max_tokens: int = 2000) -> Generator[str, None, None]:
        """"""
        try:
            payload = {
                "model": model,
                "prompt": prompt,
                "stream": True,
                "options": {
                    "temperature": temperature,
                    "num_predict": max_tokens
                }
            }
            
            if system:
                payload["system"] = system
            
            response = requests.post(
                f"{self.base_url}/api/generate",
                json=payload,
                stream=True,
                timeout=300
            )
            response.raise_for_status()
            
            for line in response.iter_lines():
                if line:
                    try:
                        data = json.loads(line)
                        if 'response' in data:
                            yield data['response']
                        if data.get('done', False):
                            break
                    except json.JSONDecodeError:
                        continue
                        
        except Exception as e:
            self.logger.error(f": {e}")
            yield f": {str(e)}"
    
    def chat(self, model: str, messages: List[Dict[str, str]], 
             temperature: float = 0.7, max_tokens: int = 2000) -> Dict[str, Any]:
        """"""
        try:
            payload = {
                "model": model,
                "messages": messages,
                "stream": False,
                "options": {
                    "temperature": temperature,
                    "num_predict": max_tokens
                }
            }
            
            response = requests.post(
                f"{self.base_url}/api/chat",
                json=payload,
                timeout=300
            )
            response.raise_for_status()
            
            return response.json()
            
        except Exception as e:
            self.logger.error(f": {e}")
            return {"error": str(e)}
    
    def _handle_stream_response(self, response) -> Dict[str, Any]:
        """"""
        full_response = ""
        
        for line in response.iter_lines():
            if line:
                try:
                    data = json.loads(line)
                    if 'response' in data:
                        full_response += data['response']
                    if data.get('done', False):
                        return {
                            "response": full_response,
                            "done": True
                        }
                except json.JSONDecodeError:
                    continue
        
        return {"response": full_response, "done": True}
    
    def get_model_info(self, model_name: str) -> Dict[str, Any]:
        """"""
        try:
            response = requests.post(
                f"{self.base_url}/api/show",
                json={"name": model_name}
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            self.logger.error(f": {e}")
            return {}
    
    def delete_model(self, model_name: str) -> bool:
        """"""
        try:
            response = requests.delete(
                f"{self.base_url}/api/delete",
                json={"name": model_name}
            )
            response.raise_for_status()
            self.logger.info(f": {model_name}")
            return True
        except Exception as e:
            self.logger.error(f": {e}")
            return False
    
    def health_check(self) -> Dict[str, Any]:
        """"""
        try:
            start_time = time.time()
            response = requests.get(f"{self.base_url}/api/tags", timeout=5)
            response_time = time.time() - start_time
            
            if response.status_code == 200:
                models = response.json().get('models', [])
                return {
                    "status": "healthy",
                    "response_time": response_time,
                    "models_count": len(models),
                    "available_models": [model['name'] for model in models]
                }
            else:
                return {
                    "status": "unhealthy",
                    "error": f"HTTP {response.status_code}"
                }
        except Exception as e:
            return {
                "status": "unhealthy",
                "error": str(e)
            }

# 
RECOMMENDED_MODELS = {
    "chinese": {
        "small": "qwen2:1.5b",
        "medium": "qwen2:7b", 
        "large": "qwen2:14b"
    },
    "english": {
        "small": "llama3.2:3b",
        "medium": "llama3.1:8b",
        "large": "llama3.1:70b"
    },
    "multilingual": {
        "small": "gemma2:2b",
        "medium": "gemma2:9b",
        "large": "gemma2:27b"
    }
}

def get_recommended_model(language: str = "chinese", size: str = "medium") -> str:
    """"""
    return RECOMMENDED_MODELS.get(language, RECOMMENDED_MODELS["chinese"]).get(size, "qwen2:7b")